import subprocess
import sys

def install_requirements():
    # 安装基本依赖
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "."])
    
    # 安装CUDA相关依赖
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", ".[cuda]"])
    
    # 安装特定版本的PyTorch
    torch_commands = [
        "torch --index-url https://download.pytorch.org/whl/cu121",
        "torchvision --index-url https://download.pytorch.org/whl/cu121",
        "torchaudio --index-url https://download.pytorch.org/whl/cu121"
    ]
    
    for cmd in torch_commands:
        subprocess.check_call([sys.executable, "-m", "pip", "install", *cmd.split()])

if __name__ == "__main__":
    install_requirements()